#include "RayTracingPipeline.h"
#include "../../Core/CoreGlobals.h"
#include "../../Utils/LogUtil.h"
#include "../../Templates/NewObject.h"
#include "../../Thread/AsyncTask.h"
#include "../../Templates/Alignmen.h"
#include "../RenderState.h"
#include "../../Game/Process/RenderProcess.h"


class CompileLibraryTask : public NonAbandonableTask
{
	friend class ManualAsyncTask<CompileLibraryTask>;

	VkDeferredOperationKHR mDeferredOperation;

	CompileLibraryTask(VkDeferredOperationKHR deferredOperation) :
		mDeferredOperation(deferredOperation)
	{}

	void DoWork()
	{
		// Wait until the thread has finished its work
		VkResult result = vkDeferredOperationJoinKHR(GVulkanRHI->GetLogicalDevice(), mDeferredOperation);
		// A return value of SUCCESS means the pipeline compilation is done.
		// THREAD_DONE indicates that thread has no work to do for this task
		// (e.g. the operation could not be split into that many threads)
		// THREAD_IDLE indicates the thread has finished its task, but the overall pipeline
		// compilation is not finished.
		// In the last two cases, more work could be performed by the thread, such as waiting
		// for another deferred operation
		assert(result == VK_SUCCESS || result == VK_THREAD_DONE_KHR || result == VK_THREAD_IDLE_KHR);
	}
};

RayTracingPipeline::Builder::Builder() noexcept = default;
RayTracingPipeline::Builder::~Builder() noexcept = default;
RayTracingPipeline::Builder::Builder(RayTracingPipeline::Builder const& rhs) noexcept = default;
RayTracingPipeline::Builder::Builder(RayTracingPipeline::Builder && rhs) noexcept = default;
RayTracingPipeline::Builder& RayTracingPipeline::Builder::operator=(RayTracingPipeline::Builder const& rhs) noexcept = default;
RayTracingPipeline::Builder& RayTracingPipeline::Builder::operator=(RayTracingPipeline::Builder && rhs) noexcept = default;

RayTracingPipeline::Builder& RayTracingPipeline::Builder::SetDescriptor(RefCountPtr<RHIDescriptor> descriptor) noexcept
{
	mImpl->descriptor = descriptor;
	return *this;
}

RayTracingPipeline::Builder& RayTracingPipeline::Builder::AddShaderStage(ShaderPtr<RHIShader> shader) noexcept
{
	// make sure shaderStage is add before shaderGroup
	assert(mImpl->shaderGroups.size() == 0);
	if (mImpl->shaderStages.size() == 0)
	{
		// make sure the first shader is raygen shader
		assert(shader->GetShaderStage() == VK_SHADER_STAGE_RAYGEN_BIT_KHR);
	}
	mImpl->shaderStages.push_back(shader);
	return *this;
}

RayTracingPipeline::Builder& RayTracingPipeline::Builder::PushShader(ShaderPtr<RHIShader> shader) noexcept
{
	switch (shader->GetShaderStage())
	{
	case VK_SHADER_STAGE_RAYGEN_BIT_KHR:
		PushRaygen(shader);
		break;
	case VK_SHADER_STAGE_MISS_BIT_KHR:
		PushMiss(shader);
		break;
	case VK_SHADER_STAGE_CLOSEST_HIT_BIT_KHR:
		PushClosestHit(shader);
		break;
	case VK_SHADER_STAGE_ANY_HIT_BIT_KHR:
		PushAnyHit(shader);
		break;
	case VK_SHADER_STAGE_INTERSECTION_BIT_KHR:
		PushIntersection(shader);
		break;
	case VK_SHADER_STAGE_CALLABLE_BIT_KHR:
		PushCallable(shader);
		break;
	default:
		assert(0);
		break;
	}
	return *this;
}

RayTracingPipeline::Builder& RayTracingPipeline::Builder::PushRaygen(ShaderPtr<RHIShader> shader) noexcept
{
	assert(shader->GetShaderStage() == VK_SHADER_STAGE_RAYGEN_BIT_KHR);
	assert(mImpl->tempCroup.generalShader == VK_SHADER_UNUSED_KHR);
	mImpl->tempCroup.generalShader = GetShaderIndex(shader);
	return *this;
}

RayTracingPipeline::Builder& RayTracingPipeline::Builder::PushMiss(ShaderPtr<RHIShader> shader) noexcept
{
	assert(shader->GetShaderStage() == VK_SHADER_STAGE_MISS_BIT_KHR);
	assert(mImpl->tempCroup.generalShader == VK_SHADER_UNUSED_KHR);
	mImpl->tempCroup.generalShader = GetShaderIndex(shader);
	return *this;
}

RayTracingPipeline::Builder& RayTracingPipeline::Builder::PushClosestHit(ShaderPtr<RHIShader> shader) noexcept
{
	assert(shader->GetShaderStage() == VK_SHADER_STAGE_CLOSEST_HIT_BIT_KHR);
	assert(mImpl->tempCroup.closestHitShader == VK_SHADER_UNUSED_KHR);
	mImpl->tempCroup.closestHitShader = GetShaderIndex(shader);
	return *this;
}

RayTracingPipeline::Builder& RayTracingPipeline::Builder::PushAnyHit(ShaderPtr<RHIShader> shader) noexcept
{
	assert(shader->GetShaderStage() == VK_SHADER_STAGE_ANY_HIT_BIT_KHR);
	assert(mImpl->tempCroup.anyHitShader == VK_SHADER_UNUSED_KHR);
	mImpl->tempCroup.anyHitShader = GetShaderIndex(shader);
	return *this;
}

RayTracingPipeline::Builder& RayTracingPipeline::Builder::PushIntersection(ShaderPtr<RHIShader> shader) noexcept
{
	assert(shader->GetShaderStage() == VK_SHADER_STAGE_INTERSECTION_BIT_KHR);
	assert(mImpl->tempCroup.intersectionShader == VK_SHADER_UNUSED_KHR);
	mImpl->tempCroup.intersectionShader = GetShaderIndex(shader);
	return *this;
}

RayTracingPipeline::Builder& RayTracingPipeline::Builder::PushCallable(ShaderPtr<RHIShader> shader) noexcept
{
	assert(shader->GetShaderStage() == VK_SHADER_STAGE_CALLABLE_BIT_KHR);
	assert(mImpl->tempCroup.generalShader == VK_SHADER_UNUSED_KHR);
	mImpl->tempCroup.generalShader = GetShaderIndex(shader);
	return *this;
}

void RayTracingPipeline::Builder::PushData(uint8 * data, uint64 dataSize)
{
	assert(mImpl->dataPtr == nullptr);
	mImpl->dataPtr = data;
	mImpl->dataSize = dataSize;
}

RayTracingPipeline::Builder& RayTracingPipeline::Builder::PackageGroup() noexcept
{
	RayTracingGroupType groupType = RayTracingGroupType::Raygen;

	if (mImpl->tempCroup.generalShader != VK_SHADER_UNUSED_KHR)
	{
		// Raygen, Miss, Callable only set at generalShader and the type are VK_RAY_TRACING_SHADER_GROUP_TYPE_GENERAL_KHR
		mImpl->tempCroup.type = VK_RAY_TRACING_SHADER_GROUP_TYPE_GENERAL_KHR;
		// Find the shader group type
		switch (mImpl->shaderStages[mImpl->tempCroup.generalShader]->GetShaderStage())
		{
		case VK_SHADER_STAGE_RAYGEN_BIT_KHR:
			groupType = RayTracingGroupType::Raygen;
			break;
		case VK_SHADER_STAGE_MISS_BIT_KHR:
			groupType = RayTracingGroupType::Miss;
			break;
		case VK_SHADER_STAGE_CALLABLE_BIT_KHR:
			groupType = RayTracingGroupType::Callable;
			break;
		default:
			assert(0);
			break;
		}
	}
	else
	{
		if (mImpl->tempCroup.intersectionShader != VK_SHADER_UNUSED_KHR)
		{
			// if intersectionShader set, type must be VK_RAY_TRACING_SHADER_GROUP_TYPE_PROCEDURAL_HIT_GROUP_KHR
			mImpl->tempCroup.type = VK_RAY_TRACING_SHADER_GROUP_TYPE_PROCEDURAL_HIT_GROUP_KHR;
		}
		else
		{
			// if closestHitShader and anyHitShader both set or only set one no matter witch, type always be VK_RAY_TRACING_SHADER_GROUP_TYPE_TRIANGLES_HIT_GROUP_KHR
			// closestHitShader and anyHitShader 两个都设置或者只设置其中一个, 都是Hit Group类型
			mImpl->tempCroup.type = VK_RAY_TRACING_SHADER_GROUP_TYPE_TRIANGLES_HIT_GROUP_KHR;
		}
		groupType = RayTracingGroupType::Hit;
	}

	if (mImpl->dataPtr != nullptr)
	{
		mImpl->dataGroups[groupType].insert({ static_cast<uint32>(mImpl->shaderGroups.size()), std::vector<uint8>(mImpl->dataPtr, mImpl->dataPtr + mImpl->dataSize) });
	}

	mImpl->shaderGroups.push_back(mImpl->tempCroup);

	// reset tempGroup
	mImpl->tempCroup.type = VK_RAY_TRACING_SHADER_GROUP_TYPE_GENERAL_KHR;
	mImpl->tempCroup.generalShader = VK_SHADER_UNUSED_KHR;
	mImpl->tempCroup.closestHitShader = VK_SHADER_UNUSED_KHR;
	mImpl->tempCroup.anyHitShader = VK_SHADER_UNUSED_KHR;
	mImpl->tempCroup.intersectionShader = VK_SHADER_UNUSED_KHR;

	mImpl->dataPtr = nullptr;
	mImpl->dataSize = 0;

	return *this;
}

RayTracingPipeline::Builder& RayTracingPipeline::Builder::SetMaxPipelineRayRecursionDepth(uint32 maxPipelineRayRecursionDepth) noexcept
{
	mImpl->maxPipelineRayRecursionDepth = maxPipelineRayRecursionDepth;
	return *this;
}

RayTracingPipeline::Builder& RayTracingPipeline::Builder::AddRayTracingLibrary(RefCountPtr<RayTracingLibrary> rayTracingLibrary) noexcept
{
	mImpl->rayTracingLibrarys.push_back(rayTracingLibrary);
	return *this;
}

RayTracingPipeline::Builder& RayTracingPipeline::Builder::AddDynamicState(VkDynamicState dynamicState) noexcept
{
	mImpl->dynamicStates.push_back(dynamicState);
	return *this;
}

RayTracingPipeline::Builder& RayTracingPipeline::Builder::SetBasePipeline(RefCountPtr<RHIPipeline> basePipeline, int32 index /*= -1*/) noexcept
{
	mImpl->basePipelineHandle = basePipeline;
	mImpl->basePipelineIndex = index;
	return *this;
}

RefCountPtr<RayTracingPipeline> RayTracingPipeline::Builder::Build()
{
	assert(IsInRenderThread());
	return RefCountPtr<RayTracingPipeline>(NewObject<RayTracingPipeline>(*mImpl));
}

RefCountPtr<RayTracingPipeline> RayTracingPipeline::Builder::Build(RenderProcess * process, const std::string & name)
{
	assert(IsInRenderThread());
	RayTracingPipeline* rayTracingPipeline = NewObject<RayTracingPipeline>(*mImpl);
	process->AddPipeline(name, rayTracingPipeline);
	return RefCountPtr<RayTracingPipeline>(rayTracingPipeline);
}

void RayTracingPipeline::ReleaseResource()
{
	if (IsInRenderThread())
	{
		if (mProcess != nullptr)
		{
			mProcess->RemovePipeline(mName, this);
		}
		DeleteObject<RayTracingPipeline>(this);
	}
	else
	{
		AsyncTask(GRenderThreadId, [this]() {
			if (mProcess != nullptr)
			{
				mProcess->RemovePipeline(mName, this);
			}
			DeleteObject<RayTracingPipeline>(this);
			});
	}
}

RayTracingPipeline::~RayTracingPipeline()
{
	for (uint32 groupType = 0; groupType < 4; ++groupType)
	{
		if (mGroupBuffers[groupType] != VK_NULL_HANDLE)
		{
			// vkDestroyBuffer(GVulkanRHI->GetLogicalDevice(), mGroupBuffers[groupType], nullptr);
			GRenderState->DelayGCBuffer(mGroupBuffers[groupType]);
		}
		if (mGroupMemorys[groupType] != VK_NULL_HANDLE)
		{
			// vkFreeMemory(GVulkanRHI->GetLogicalDevice(), mGroupMemorys[groupType], nullptr);
			GRenderState->DelayGCDeviceMemory(mGroupMemorys[groupType]);
		}
	}
}

std::array<VkStridedDeviceAddressRegionKHR, 4> RayTracingPipeline::GetRegions() const
{
	return mGroupRegions;
}

const VkStridedDeviceAddressRegionKHR& RayTracingPipeline::GetRaygenRegion() const
{
	return mGroupRegions[RayTracingGroupType::Raygen];
}

const VkStridedDeviceAddressRegionKHR& RayTracingPipeline::GetMissRegion() const
{
	return mGroupRegions[RayTracingGroupType::Miss];
}

const VkStridedDeviceAddressRegionKHR& RayTracingPipeline::GetHitRegion() const
{
	return mGroupRegions[RayTracingGroupType::Hit];
}

const VkStridedDeviceAddressRegionKHR& RayTracingPipeline::GetCallableRegion() const
{
	return mGroupRegions[RayTracingGroupType::Callable];
}

inline std::vector<ShaderPtr<RHIShader>>& RayTracingPipeline::GetShaderStages()
{
	return mShaderStages;
}

uint32 RayTracingPipeline::Builder::GetShaderIndex(ShaderPtr<RHIShader> shader) const
{
	for (uint32 i = 0; i < mImpl->shaderStages.size(); ++i)
	{
		if (mImpl->shaderStages[i] == shader)
		{
			return i;
		}
	}
	// no shader exist
	// make sure the shader is add to shaderStages before make shaderGroup
	assert(0);
	return 0;
}

RayTracingPipeline::RayTracingPipeline(BuilderDetails & builderDetails)
{
	mDescriptor = builderDetails.descriptor;

	assert(builderDetails.maxPipelineRayRecursionDepth > 0 && builderDetails.maxPipelineRayRecursionDepth <= GVulkanRHI->GetRayTracingPipelineProperties().maxRayRecursionDepth);


	mShaderStages.assign(builderDetails.shaderStages.begin(), builderDetails.shaderStages.end());
	mShaderGroups.assign(builderDetails.shaderGroups.begin(), builderDetails.shaderGroups.end());
	for (uint32 i = 0; i < 4; ++i)
	{
		for (const auto& data : builderDetails.dataGroups[i])
		{
			mDataGroups[i].insert(data);
		}
	}

	mRayTracingLibrarys.assign(builderDetails.rayTracingLibrarys.begin(), builderDetails.rayTracingLibrarys.end());

	mMaxPipelineRayRecursionDepth = builderDetails.maxPipelineRayRecursionDepth;

	VkRayTracingPipelineCreateInfoKHR rayTracingPipelineCreateInfo{};
	rayTracingPipelineCreateInfo.sType = VK_STRUCTURE_TYPE_RAY_TRACING_PIPELINE_CREATE_INFO_KHR;

	VkPipelineLibraryCreateInfoKHR pipelineLibraryCreateInfo{};
	std::vector<VkPipeline> libraryPipelines;
	VkRayTracingPipelineInterfaceCreateInfoKHR libraryInterface;
	if (mRayTracingLibrarys.size() > 0)
	{
		libraryInterface = mRayTracingLibrarys[0]->GetRayTracingPipelineInterface();
		for (auto& rayTracingLibrary : mRayTracingLibrarys)
		{
			libraryPipelines.push_back(rayTracingLibrary->GetPipeline());

			// target pipeline and library must have the same libraryInterface
			assert(libraryInterface.maxPipelineRayHitAttributeSize == rayTracingLibrary->GetRayTracingPipelineInterface().maxPipelineRayHitAttributeSize);
			assert(libraryInterface.maxPipelineRayPayloadSize == rayTracingLibrary->GetRayTracingPipelineInterface().maxPipelineRayPayloadSize);

			// target pipeline and library must have the same mMaxPipelineRayRecursionDepth
			assert(mMaxPipelineRayRecursionDepth == rayTracingLibrary->GetMaxPipelineRayRecursionDepth());

			// TODO target pipeline and library must have the same layout

		}
		pipelineLibraryCreateInfo.sType = VK_STRUCTURE_TYPE_PIPELINE_LIBRARY_CREATE_INFO_KHR;
		pipelineLibraryCreateInfo.libraryCount = static_cast<uint32>(mRayTracingLibrarys.size());
		pipelineLibraryCreateInfo.pLibraries = libraryPipelines.data();

		rayTracingPipelineCreateInfo.pLibraryInfo = &pipelineLibraryCreateInfo;
		rayTracingPipelineCreateInfo.pLibraryInterface = &libraryInterface;

		rayTracingPipelineCreateInfo.maxPipelineRayRecursionDepth = mMaxPipelineRayRecursionDepth;
	}
	else
	{
		rayTracingPipelineCreateInfo.maxPipelineRayRecursionDepth = mMaxPipelineRayRecursionDepth;
	}

	rayTracingPipelineCreateInfo.layout = mDescriptor->GetPipelineLayout();

	std::vector<VkPipelineShaderStageCreateInfo> shadeStages;
	for (auto& shader : mShaderStages)
		shadeStages.push_back(shader->GetShaderStageCreateInfo());

	rayTracingPipelineCreateInfo.stageCount = static_cast<uint32>(shadeStages.size());
	rayTracingPipelineCreateInfo.pStages = shadeStages.data();
	rayTracingPipelineCreateInfo.groupCount = static_cast<uint32>(mShaderGroups.size());
	rayTracingPipelineCreateInfo.pGroups = mShaderGroups.data();


	VkPipelineDynamicStateCreateInfo dynamicStateCreateInfo{};
	if (builderDetails.dynamicStates.size() > 0)
	{
		dynamicStateCreateInfo.sType = VK_STRUCTURE_TYPE_PIPELINE_DYNAMIC_STATE_CREATE_INFO;
		dynamicStateCreateInfo.dynamicStateCount = static_cast<uint32>(builderDetails.dynamicStates.size());
		dynamicStateCreateInfo.pDynamicStates = builderDetails.dynamicStates.data();
		rayTracingPipelineCreateInfo.pDynamicState = &dynamicStateCreateInfo;
	}
	else
	{
		rayTracingPipelineCreateInfo.pDynamicState = nullptr;
	}

	if (builderDetails.basePipelineHandle.IsValid())
	{
		// 在基础管道之后创建的所有管道都将是衍生品
		rayTracingPipelineCreateInfo.flags = VK_PIPELINE_CREATE_DERIVATIVE_BIT;
		// 基础管道将是我们创建的第一个管道
		rayTracingPipelineCreateInfo.basePipelineHandle = builderDetails.basePipelineHandle->GetPipeline();
		// 只允许对基本管道使用句柄或索引
		// 当我们使用句柄时，我们必须将索引设置为 -1（参见规范的第 9.5 节）
		rayTracingPipelineCreateInfo.basePipelineIndex = builderDetails.basePipelineIndex;
	}
	else
	{
		rayTracingPipelineCreateInfo.basePipelineHandle = VK_NULL_HANDLE;
		rayTracingPipelineCreateInfo.basePipelineIndex = -1;
	}

	if (mRayTracingLibrarys.size() > 0)
	{
		// Use multi thread to create pipeline
		// Deferred operations allow the driver to parallelize the pipeline compilation on several threads
		// Create a deferred operation
		VkDeferredOperationKHR deferredOperation;
		VK_CHECK_RESULT(vkCreateDeferredOperationKHR(GVulkanRHI->GetLogicalDevice(), nullptr, &deferredOperation));

		// The pipeline creation is called with the deferred operation. Instead of blocking until
		// the compilation is done, the call returns immediately
		vkCreateRayTracingPipelinesKHR(GVulkanRHI->GetLogicalDevice(), deferredOperation, {}, 1, &rayTracingPipelineCreateInfo, nullptr, &mPipeline);

#if 1
		std::vector<ManualAsyncTask<CompileLibraryTask>*> compileTasks(GThreadPool->GetNumThreads());
		for (uint32 i = 0; i < compileTasks.size(); ++i)
		{
			compileTasks[i] = new ManualAsyncTask<CompileLibraryTask>(deferredOperation);
			compileTasks[i]->Start();
		}
		for (uint32 i = 0; i < compileTasks.size(); ++i)
		{
			compileTasks[i]->EnsureCompletion();
			delete compileTasks[i];
			compileTasks[i] = nullptr;
		}
#else
		uint32 maxThreads = 8;
		uint32 threadCount = glm::min(vkGetDeferredOperationMaxConcurrencyKHR(GVulkanRHI->GetLogicalDevice(), deferredOperation), maxThreads);
		std::vector<std::future<void>> joins;
		for (uint32_t i = 0; i < threadCount; i++)
		{
			VkDevice device{ GVulkanRHI->GetLogicalDevice() };
			joins.emplace_back(std::async(std::launch::async, [device, deferredOperation]() {
				// Wait until the thread has finished its work
				VkResult result = vkDeferredOperationJoinKHR(device, deferredOperation);
				// A return value of SUCCESS means the pipeline compilation is done.
				// THREAD_DONE indicates that thread has no work to do for this task
				// (e.g. the operation could not be split into that many threads)
				// THREAD_IDLE indicates the thread has finished its task, but the overall pipeline
				// compilation is not finished.
				// In the last two cases, more work could be performed by the thread, such as waiting
				// for another deferred operation
				assert(result == VK_SUCCESS || result == VK_THREAD_DONE_KHR || result == VK_THREAD_IDLE_KHR);
				}));
		}
		// Wait for all threads to finish
		for (auto& f : joins)
		{
			f.get();
		}
#endif
		// Once the deferred operation is complete, check for compilation success
		VK_CHECK_RESULT(vkGetDeferredOperationResultKHR(GVulkanRHI->GetLogicalDevice(), deferredOperation));
		// Destroy the deferred operation
		vkDestroyDeferredOperationKHR(GVulkanRHI->GetLogicalDevice(), deferredOperation, nullptr);
	}
	else
	{
		// no library , create pipeline direct
		VK_CHECK_RESULT(vkCreateRayTracingPipelinesKHR(GVulkanRHI->GetLogicalDevice(), VK_NULL_HANDLE, VK_NULL_HANDLE, 1, &rayTracingPipelineCreateInfo, nullptr, &mPipeline));
	}

	CreateShaderBingingTable();
}

void RayTracingPipeline::CreateShaderBingingTable()
{
	const uint32 handleSize = GVulkanRHI->GetRayTracingPipelineProperties().shaderGroupHandleSize;
	const uint32 handleAlignment = GVulkanRHI->GetRayTracingPipelineProperties().shaderGroupHandleAlignment;

	// Get the total number of groups and handle index position
	uint32 totalGroupCount = 0;

	// save the group index in group index
	std::array<std::vector<uint32>, 4> groupIndexs;

	// Local data to groups (Shader Record), use the same index of groupIndexs
	std::array<std::unordered_map<uint32, std::vector<uint8>>, 4> groupDatas;

	// Hold the stage buffer for group type shader buffer
	std::array<VkBuffer, 4> stagingBuffers;
	std::array<VkDeviceMemory, 4> stagingMemorys;

	// init everythings
	for (uint32 i = 0; i < 4; ++i)
	{
		groupIndexs[i] = {};

		stagingBuffers[i] = VK_NULL_HANDLE;
		stagingMemorys[i] = VK_NULL_HANDLE;
		mGroupBuffers[i] = VK_NULL_HANDLE;
		mGroupMemorys[i] = VK_NULL_HANDLE;

		mGroupRegions[i] = {};
	}

	uint32 groupOffset = 0;
	for (uint32 i = 0; i < mRayTracingLibrarys.size() + 1; ++i)
	{
		// first one for target pipeline
		std::vector<VkRayTracingShaderGroupCreateInfoKHR>& shaderGroups = i == 0 ? mShaderGroups : mRayTracingLibrarys[i - 1]->GetShaderGroups();
		std::vector<ShaderPtr<RHIShader>>& shaderStages = i == 0 ? mShaderStages : mRayTracingLibrarys[i - 1]->GetShaderStages();
		std::array<std::unordered_map<uint32, std::vector<uint8>>, 4>& dataGroups = i == 0 ? mDataGroups : mRayTracingLibrarys[i - 1]->GetDataGroups();

		totalGroupCount += static_cast<uint32>(shaderGroups.size());

		// Finding the handle position of each group, splitting by raygen, missand hit group
		for (uint32 g = 0; g < shaderGroups.size(); ++g)
		{
			RayTracingGroupType groupType = RayTracingGroupType::Raygen;
			if (shaderGroups[g].type == VK_RAY_TRACING_SHADER_GROUP_TYPE_GENERAL_KHR)
			{
				switch (shaderStages[shaderGroups[g].generalShader]->GetShaderStage())
				{
				case VK_SHADER_STAGE_RAYGEN_BIT_KHR:
					groupType = RayTracingGroupType::Raygen;
					break;
				case VK_SHADER_STAGE_MISS_BIT_KHR:
					groupType = RayTracingGroupType::Miss;
					break;
				case VK_SHADER_STAGE_CALLABLE_BIT_KHR:
					groupType = RayTracingGroupType::Callable;
					break;
				default:
					assert(0);
					break;
				}
			}
			else
			{
				groupType = RayTracingGroupType::Hit;
			}

			// insert group index
			groupIndexs[groupType].push_back(g + groupOffset);

			// if the data of g is exist, put it in groupDatas by g + groupOffset
			if (dataGroups[groupType].find(g) != dataGroups[groupType].end())
			{
				groupDatas[groupType].insert({ g + groupOffset, dataGroups[groupType].find(g)->second });
			}
		}
		groupOffset += shaderGroups.size();
	}

	// Fetch all the shader handles used in the pipeline, so that they can be written in the SBT
	uint32 sbtSize = totalGroupCount * handleSize;
	std::vector<uint8> shaderHandleStorage(sbtSize);

	VK_CHECK_RESULT(vkGetRayTracingShaderGroupHandlesKHR(GVulkanRHI->GetLogicalDevice(), mPipeline, 0, totalGroupCount, sbtSize, shaderHandleStorage.data()));

	// Find the max stride, minimum is the handle size + size of 'data (if any)' aligned to shaderGroupBaseAlignment
	std::function<void(std::unordered_map<uint32, std::vector<uint8>>&, VkStridedDeviceAddressRegionKHR&, uint32 groupCount)> FindStride = [handleSize, handleAlignment](std::unordered_map<uint32, std::vector<uint8>>& datas, VkStridedDeviceAddressRegionKHR& region, uint32 groupCount)
	{
		region.stride = Align(handleSize, handleAlignment);
		for (auto& data : datas)
		{
			VkDeviceSize needSize = Align(static_cast<uint32>(handleSize + data.second.size()), handleAlignment);
			region.stride = glm::max(region.stride, needSize);
		}
		region.size = static_cast<VkDeviceSize>(region.stride * groupCount);
	};

	for (uint32 groupType = 0; groupType < 4; ++groupType)
	{
		FindStride(groupDatas[groupType], mGroupRegions[groupType], static_cast<uint32>(groupIndexs[groupType].size()));
	}

	std::function<void(uint8*, std::vector<uint32>&, std::unordered_map<uint32, std::vector<uint8>>&, VkDeviceSize)> CopyHandle = [&shaderHandleStorage, handleSize](uint8* maped, std::vector<uint32>& indices, std::unordered_map<uint32, std::vector<uint8>>& datas, VkDeviceSize stride)
	{
		uint8* buffer = maped;
		for (auto& index : indices)
		{
			uint8* temp = buffer;
			// Copy the handle
			memcpy(buffer, shaderHandleStorage.data() + (index * handleSize), handleSize);
			// if data exist, copy data next to the handle
			std::unordered_map<uint32, std::vector<uint8>>::iterator it = datas.find(index);
			if (it != datas.end())
			{
				buffer += handleSize;
				memcpy(buffer, it->second.data(), it->second.size() * sizeof(uint8));
			}
			buffer = temp + stride;
		}
	};

	for (uint32 groupType = 0; groupType < 4; ++groupType)
	{
		// same group no data, should be {}
		if (mGroupRegions[groupType].size == 0)
		{
			mGroupRegions[groupType] = {};
			continue;
		}

		uint8* maped = nullptr;
		maped = CreateStagingBuffer(mGroupRegions[groupType].size, stagingBuffers[groupType], stagingMemorys[groupType]);
		CopyHandle(maped, groupIndexs[groupType], groupDatas[groupType], mGroupRegions[groupType].stride);
		vkUnmapMemory(GVulkanRHI->GetLogicalDevice(), stagingMemorys[groupType]);
		mGroupRegions[groupType].deviceAddress = CreateGroupBuffer(mGroupRegions[groupType].size, mGroupBuffers[groupType], mGroupMemorys[groupType]);
		CopyBuffer(stagingBuffers[groupType], mGroupBuffers[groupType], mGroupRegions[groupType].size);
		vkDestroyBuffer(GVulkanRHI->GetLogicalDevice(), stagingBuffers[groupType], nullptr);
		vkFreeMemory(GVulkanRHI->GetLogicalDevice(), stagingMemorys[groupType], nullptr);
	}
}

uint8* RayTracingPipeline::CreateStagingBuffer(VkDeviceSize size, VkBuffer & buffer, VkDeviceMemory & bufferMemory)
{
	VkBufferCreateInfo bufferCreateInfo{};
	bufferCreateInfo.sType = VK_STRUCTURE_TYPE_BUFFER_CREATE_INFO;
	bufferCreateInfo.size = size;
	bufferCreateInfo.usage = VK_BUFFER_USAGE_TRANSFER_SRC_BIT;
	bufferCreateInfo.sharingMode = VK_SHARING_MODE_EXCLUSIVE;

	VK_CHECK_RESULT(vkCreateBuffer(GVulkanRHI->GetLogicalDevice(), &bufferCreateInfo, nullptr, &buffer));

	VkMemoryRequirements memRequirements;
	vkGetBufferMemoryRequirements(GVulkanRHI->GetLogicalDevice(), buffer, &memRequirements);

	VkMemoryAllocateInfo memoryAllocateInfo{};
	memoryAllocateInfo.sType = VK_STRUCTURE_TYPE_MEMORY_ALLOCATE_INFO;
	memoryAllocateInfo.allocationSize = memRequirements.size;

	memoryAllocateInfo.memoryTypeIndex = GVulkanRHI->FindMemoryType(memRequirements.memoryTypeBits, VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT | VK_MEMORY_PROPERTY_HOST_COHERENT_BIT);

	VK_CHECK_RESULT(vkAllocateMemory(GVulkanRHI->GetLogicalDevice(), &memoryAllocateInfo, nullptr, &bufferMemory));

	vkBindBufferMemory(GVulkanRHI->GetLogicalDevice(), buffer, bufferMemory, 0);

	void* maped;
	vkMapMemory(GVulkanRHI->GetLogicalDevice(), bufferMemory, 0, size, 0, &maped);
	return static_cast<uint8*>(maped);
}

VkDeviceAddress RayTracingPipeline::CreateGroupBuffer(VkDeviceSize size, VkBuffer & buffer, VkDeviceMemory & bufferMemory)
{
	VkBufferCreateInfo bufferCreateInfo{};
	bufferCreateInfo.sType = VK_STRUCTURE_TYPE_BUFFER_CREATE_INFO;
	bufferCreateInfo.size = size;
	bufferCreateInfo.usage = VK_BUFFER_USAGE_TRANSFER_DST_BIT | VK_BUFFER_USAGE_SHADER_DEVICE_ADDRESS_BIT;
	bufferCreateInfo.sharingMode = VK_SHARING_MODE_EXCLUSIVE;

	VK_CHECK_RESULT(vkCreateBuffer(GVulkanRHI->GetLogicalDevice(), &bufferCreateInfo, nullptr, &buffer));

	VkMemoryRequirements memRequirements;
	vkGetBufferMemoryRequirements(GVulkanRHI->GetLogicalDevice(), buffer, &memRequirements);

	VkMemoryAllocateInfo memoryAllocateInfo{};
	memoryAllocateInfo.sType = VK_STRUCTURE_TYPE_MEMORY_ALLOCATE_INFO;
	memoryAllocateInfo.allocationSize = memRequirements.size;

	memoryAllocateInfo.memoryTypeIndex = GVulkanRHI->FindMemoryType(memRequirements.memoryTypeBits, VK_MEMORY_PROPERTY_DEVICE_LOCAL_BIT);

	// If usage & VK_BUFFER_USAGE_SHADER_DEVICE_ADDRESS_BIT, need to add a VkMemoryAllocateFlagsInfo to pNext of VkMemoryAllocateInfo
	VkMemoryAllocateFlagsInfo memoryAllocateFlagsInfo{};
	memoryAllocateFlagsInfo.sType = VK_STRUCTURE_TYPE_MEMORY_ALLOCATE_FLAGS_INFO;
	memoryAllocateFlagsInfo.flags = VK_MEMORY_ALLOCATE_DEVICE_ADDRESS_BIT;
	memoryAllocateInfo.pNext = &memoryAllocateFlagsInfo;

	VK_CHECK_RESULT(vkAllocateMemory(GVulkanRHI->GetLogicalDevice(), &memoryAllocateInfo, nullptr, &bufferMemory));

	vkBindBufferMemory(GVulkanRHI->GetLogicalDevice(), buffer, bufferMemory, 0);

	VkBufferDeviceAddressInfoKHR bufferDeviceAddressInfo{};
	bufferDeviceAddressInfo.sType = VK_STRUCTURE_TYPE_BUFFER_DEVICE_ADDRESS_INFO;
	bufferDeviceAddressInfo.buffer = buffer;
	VkDeviceAddress deviceAddress = vkGetBufferDeviceAddressKHR(GVulkanRHI->GetLogicalDevice(), &bufferDeviceAddressInfo);
	assert(deviceAddress != 0);
	return deviceAddress;
}

void RayTracingPipeline::CopyBuffer(VkBuffer srcBuffer, VkBuffer dstBuffer, VkDeviceSize size)
{
	VkCommandBuffer commandBuffer = GRenderState->BeginSingleCommand();

	VkBufferCopy copyRegion{};
	copyRegion.srcOffset = 0;
	copyRegion.dstOffset = 0;
	copyRegion.size = size;
	vkCmdCopyBuffer(commandBuffer, srcBuffer, dstBuffer, 1, &copyRegion);

	GRenderState->FlushSingleCommand(commandBuffer);
}
